{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/pj20/miniconda3/envs/txgnn_env/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO: Pandarallel will run on 64 workers.\n",
      "INFO: Pandarallel will use Memory file system to transfer data between the main process and workers.\n",
      "finish basic patient information parsing : 1.295158863067627s\n",
      "finish parsing DIAGNOSES_ICD : 1.5444056987762451s\n",
      "finish parsing PROCEDURES_ICD : 2.137986421585083s\n",
      "finish parsing PRESCRIPTIONS : 9.635247707366943s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mapping codes: 100%|██████████| 1000/1000 [00:01<00:00, 560.06it/s]\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.datasets import MIMIC3Dataset\n",
    "from pyhealth.datasets import split_by_patient, get_dataloader\n",
    "from pyhealth.models import Transformer, RNN, RETAIN, MLP\n",
    "from pyhealth.tasks import mortality_prediction_mimic3_fn, readmission_prediction_mimic3_fn, drug_recommendation_mimic3_fn, length_of_stay_prediction_mimic3_fn\n",
    "from pyhealth.trainer import Trainer\n",
    "\n",
    "dataset = MIMIC3Dataset(\n",
    "    root='/data/physionet.org/files/mimiciii/1.4/',\n",
    "    tables=[\"DIAGNOSES_ICD\", \"PROCEDURES_ICD\", \"PRESCRIPTIONS\"],\n",
    "    code_mapping={\n",
    "        # \"ICD9CM\": \"CCSCM\", \n",
    "        # \"ICD9PROC\": \"CCSPROC\",\n",
    "        \"NDC\": (\"ATC\", {\"target_kwargs\": {\"level\": 3}})\n",
    "        },\n",
    "    dev=True,\n",
    "    refresh_cache=True\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating samples for readmission_prediction_mimic3_fn: 100%|██████████| 1000/1000 [00:00<00:00, 51093.97it/s]\n"
     ]
    }
   ],
   "source": [
    "mimic3_ds = dataset.set_task(readmission_prediction_mimic3_fn)\n",
    "\n",
    "train_dataset, val_dataset, test_dataset = split_by_patient(\n",
    "    mimic3_ds, [0.8, 0.1, 0.1]\n",
    ")\n",
    "train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)\n",
    "val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)\n",
    "test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading pretrained embedding for conditions...\n",
      "Loading pretrained embedding for procedures...\n",
      "Loading pretrained embedding for drugs...\n"
     ]
    }
   ],
   "source": [
    "model_w_pre = Transformer(\n",
    "    dataset=mimic3_ds,\n",
    "    feature_keys=[\"conditions\", \"procedures\", \"drugs\"],\n",
    "    label_key=\"label\",\n",
    "    mode=\"binary\",\n",
    "    pretrained_emb=\"KG/transe\",\n",
    "    embedding_dim=256,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer(\n",
      "  (embeddings): ModuleDict(\n",
      "    (conditions): Embedding(4031, 512, padding_idx=0)\n",
      "    (procedures): Embedding(1276, 512, padding_idx=0)\n",
      "    (drugs): Embedding(194, 512, padding_idx=0)\n",
      "  )\n",
      "  (linear_layers): ModuleDict(\n",
      "    (conditions): Linear(in_features=512, out_features=256, bias=True)\n",
      "    (procedures): Linear(in_features=512, out_features=256, bias=True)\n",
      "    (drugs): Linear(in_features=512, out_features=256, bias=True)\n",
      "  )\n",
      "  (transformer): ModuleDict(\n",
      "    (conditions): TransformerLayer(\n",
      "      (transformer): ModuleList(\n",
      "        (0): TransformerBlock(\n",
      "          (attention): MultiHeadedAttention(\n",
      "            (linear_layers): ModuleList(\n",
      "              (0): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (1): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (2): Linear(in_features=256, out_features=256, bias=False)\n",
      "            )\n",
      "            (output_linear): Linear(in_features=256, out_features=256, bias=False)\n",
      "            (attention): Attention()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (feed_forward): PositionwiseFeedForward(\n",
      "            (w_1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "            (w_2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "            (activation): GELU(approximate=none)\n",
      "          )\n",
      "          (input_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (output_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (dropout): Dropout(p=0.5, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (procedures): TransformerLayer(\n",
      "      (transformer): ModuleList(\n",
      "        (0): TransformerBlock(\n",
      "          (attention): MultiHeadedAttention(\n",
      "            (linear_layers): ModuleList(\n",
      "              (0): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (1): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (2): Linear(in_features=256, out_features=256, bias=False)\n",
      "            )\n",
      "            (output_linear): Linear(in_features=256, out_features=256, bias=False)\n",
      "            (attention): Attention()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (feed_forward): PositionwiseFeedForward(\n",
      "            (w_1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "            (w_2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "            (activation): GELU(approximate=none)\n",
      "          )\n",
      "          (input_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (output_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (dropout): Dropout(p=0.5, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (drugs): TransformerLayer(\n",
      "      (transformer): ModuleList(\n",
      "        (0): TransformerBlock(\n",
      "          (attention): MultiHeadedAttention(\n",
      "            (linear_layers): ModuleList(\n",
      "              (0): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (1): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (2): Linear(in_features=256, out_features=256, bias=False)\n",
      "            )\n",
      "            (output_linear): Linear(in_features=256, out_features=256, bias=False)\n",
      "            (attention): Attention()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (feed_forward): PositionwiseFeedForward(\n",
      "            (w_1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "            (w_2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "            (activation): GELU(approximate=none)\n",
      "          )\n",
      "          (input_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (output_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (dropout): Dropout(p=0.5, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (fc): Linear(in_features=768, out_features=1, bias=True)\n",
      ")\n",
      "Metrics: None\n",
      "Device: cuda\n",
      "\n",
      "Training:\n",
      "Batch size: 32\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.0001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: None\n",
      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x7ff5195e7520>\n",
      "Monitor: pr_auc\n",
      "Monitor criterion: max\n",
      "Epochs: 15\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 / 15: 100%|██████████| 243/243 [00:02<00:00, 81.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-0, step-243 ---\n",
      "loss: 2.6863\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 259.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-0, step-243 ---\n",
      "pr_auc: 0.5940\n",
      "roc_auc: 0.5701\n",
      "f1: 0.6520\n",
      "loss: 0.7162\n",
      "New best pr_auc score (0.5940) at epoch-0, step-243\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1 / 15: 100%|██████████| 243/243 [00:02<00:00, 86.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-1, step-486 ---\n",
      "loss: 1.3624\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 262.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-1, step-486 ---\n",
      "pr_auc: 0.6169\n",
      "roc_auc: 0.5924\n",
      "f1: 0.6379\n",
      "loss: 0.7208\n",
      "New best pr_auc score (0.6169) at epoch-1, step-486\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2 / 15: 100%|██████████| 243/243 [00:02<00:00, 86.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-2, step-729 ---\n",
      "loss: 1.0040\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 255.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-2, step-729 ---\n",
      "pr_auc: 0.5691\n",
      "roc_auc: 0.5596\n",
      "f1: 0.6870\n",
      "loss: 1.0541\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3 / 15: 100%|██████████| 243/243 [00:02<00:00, 87.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-3, step-972 ---\n",
      "loss: 0.9451\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 263.53it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-3, step-972 ---\n",
      "pr_auc: 0.6197\n",
      "roc_auc: 0.5927\n",
      "f1: 0.6518\n",
      "loss: 0.6995\n",
      "New best pr_auc score (0.6197) at epoch-3, step-972\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 4 / 15: 100%|██████████| 243/243 [00:02<00:00, 87.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-4, step-1215 ---\n",
      "loss: 0.9677\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 263.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-4, step-1215 ---\n",
      "pr_auc: 0.6192\n",
      "roc_auc: 0.6033\n",
      "f1: 0.4757\n",
      "loss: 0.7513\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 5 / 15: 100%|██████████| 243/243 [00:02<00:00, 87.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-5, step-1458 ---\n",
      "loss: 0.9120\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 257.26it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-5, step-1458 ---\n",
      "pr_auc: 0.6010\n",
      "roc_auc: 0.5879\n",
      "f1: 0.1362\n",
      "loss: 0.9770\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 6 / 15: 100%|██████████| 243/243 [00:02<00:00, 85.58it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-6, step-1701 ---\n",
      "loss: 0.8684\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 166.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-6, step-1701 ---\n",
      "pr_auc: 0.6192\n",
      "roc_auc: 0.6064\n",
      "f1: 0.6898\n",
      "loss: 0.7518\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 7 / 15: 100%|██████████| 243/243 [00:02<00:00, 83.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-7, step-1944 ---\n",
      "loss: 0.8355\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 149.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-7, step-1944 ---\n",
      "pr_auc: 0.6051\n",
      "roc_auc: 0.5946\n",
      "f1: 0.1900\n",
      "loss: 0.7860\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 8 / 15: 100%|██████████| 243/243 [00:03<00:00, 72.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-8, step-2187 ---\n",
      "loss: 0.8441\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 260.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-8, step-2187 ---\n",
      "pr_auc: 0.6311\n",
      "roc_auc: 0.6174\n",
      "f1: 0.6282\n",
      "loss: 0.6806\n",
      "New best pr_auc score (0.6311) at epoch-8, step-2187\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 9 / 15: 100%|██████████| 243/243 [00:03<00:00, 71.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-9, step-2430 ---\n",
      "loss: 0.8071\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 256.04it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-9, step-2430 ---\n",
      "pr_auc: 0.6187\n",
      "roc_auc: 0.6039\n",
      "f1: 0.4269\n",
      "loss: 0.6895\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 10 / 15: 100%|██████████| 243/243 [00:03<00:00, 79.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-10, step-2673 ---\n",
      "loss: 0.7709\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 134.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-10, step-2673 ---\n",
      "pr_auc: 0.6393\n",
      "roc_auc: 0.6198\n",
      "f1: 0.6864\n",
      "loss: 0.7150\n",
      "New best pr_auc score (0.6393) at epoch-10, step-2673\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 11 / 15: 100%|██████████| 243/243 [00:03<00:00, 69.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-11, step-2916 ---\n",
      "loss: 0.7761\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 257.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-11, step-2916 ---\n",
      "pr_auc: 0.6172\n",
      "roc_auc: 0.6027\n",
      "f1: 0.3242\n",
      "loss: 0.7546\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 12 / 15: 100%|██████████| 243/243 [00:03<00:00, 78.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-12, step-3159 ---\n",
      "loss: 0.7396\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 183.00it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-12, step-3159 ---\n",
      "pr_auc: 0.6332\n",
      "roc_auc: 0.6176\n",
      "f1: 0.4724\n",
      "loss: 0.7049\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 13 / 15: 100%|██████████| 243/243 [00:03<00:00, 80.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-13, step-3402 ---\n",
      "loss: 0.7206\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 253.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-13, step-3402 ---\n",
      "pr_auc: 0.6470\n",
      "roc_auc: 0.6315\n",
      "f1: 0.6496\n",
      "loss: 0.6777\n",
      "New best pr_auc score (0.6470) at epoch-13, step-3402\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 14 / 15: 100%|██████████| 243/243 [00:03<00:00, 74.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-14, step-3645 ---\n",
      "loss: 0.6901\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 257.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-14, step-3645 ---\n",
      "pr_auc: 0.6423\n",
      "roc_auc: 0.6159\n",
      "f1: 0.6416\n",
      "loss: 0.6913\n",
      "Loaded best model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 32/32 [00:00<00:00, 256.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'pr_auc': 0.7206943820433741, 'roc_auc': 0.6867013897339791, 'f1': 0.6968174204355109, 'loss': 0.6363492039963603}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# STEP 4: define trainer\n",
    "trainer = Trainer(model=model_w_pre)\n",
    "trainer.train(\n",
    "    train_dataloader=train_dataloader,\n",
    "    val_dataloader=val_dataloader,\n",
    "    epochs=15,\n",
    "    optimizer_params = {\"lr\": 1e-4},\n",
    "    monitor=\"pr_auc\",\n",
    ")\n",
    "\n",
    "# STEP 5: evaluate\n",
    "print(trainer.evaluate(test_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_no_pre = Transformer(\n",
    "    dataset=mimic3_ds,\n",
    "    feature_keys=[\"conditions\", \"procedures\", \"drugs\"],\n",
    "    label_key=\"label\",\n",
    "    mode=\"binary\",\n",
    "    embedding_dim=256,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer(\n",
      "  (embeddings): ModuleDict(\n",
      "    (conditions): Embedding(4031, 256, padding_idx=0)\n",
      "    (procedures): Embedding(1276, 256, padding_idx=0)\n",
      "    (drugs): Embedding(194, 256, padding_idx=0)\n",
      "  )\n",
      "  (linear_layers): ModuleDict()\n",
      "  (transformer): ModuleDict(\n",
      "    (conditions): TransformerLayer(\n",
      "      (transformer): ModuleList(\n",
      "        (0): TransformerBlock(\n",
      "          (attention): MultiHeadedAttention(\n",
      "            (linear_layers): ModuleList(\n",
      "              (0): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (1): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (2): Linear(in_features=256, out_features=256, bias=False)\n",
      "            )\n",
      "            (output_linear): Linear(in_features=256, out_features=256, bias=False)\n",
      "            (attention): Attention()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (feed_forward): PositionwiseFeedForward(\n",
      "            (w_1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "            (w_2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "            (activation): GELU(approximate=none)\n",
      "          )\n",
      "          (input_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (output_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (dropout): Dropout(p=0.5, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (procedures): TransformerLayer(\n",
      "      (transformer): ModuleList(\n",
      "        (0): TransformerBlock(\n",
      "          (attention): MultiHeadedAttention(\n",
      "            (linear_layers): ModuleList(\n",
      "              (0): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (1): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (2): Linear(in_features=256, out_features=256, bias=False)\n",
      "            )\n",
      "            (output_linear): Linear(in_features=256, out_features=256, bias=False)\n",
      "            (attention): Attention()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (feed_forward): PositionwiseFeedForward(\n",
      "            (w_1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "            (w_2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "            (activation): GELU(approximate=none)\n",
      "          )\n",
      "          (input_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (output_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (dropout): Dropout(p=0.5, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (drugs): TransformerLayer(\n",
      "      (transformer): ModuleList(\n",
      "        (0): TransformerBlock(\n",
      "          (attention): MultiHeadedAttention(\n",
      "            (linear_layers): ModuleList(\n",
      "              (0): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (1): Linear(in_features=256, out_features=256, bias=False)\n",
      "              (2): Linear(in_features=256, out_features=256, bias=False)\n",
      "            )\n",
      "            (output_linear): Linear(in_features=256, out_features=256, bias=False)\n",
      "            (attention): Attention()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (feed_forward): PositionwiseFeedForward(\n",
      "            (w_1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "            (w_2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "            (activation): GELU(approximate=none)\n",
      "          )\n",
      "          (input_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (output_sublayer): SublayerConnection(\n",
      "            (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.5, inplace=False)\n",
      "          )\n",
      "          (dropout): Dropout(p=0.5, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (fc): Linear(in_features=768, out_features=1, bias=True)\n",
      ")\n",
      "Metrics: None\n",
      "Device: cuda\n",
      "\n",
      "Training:\n",
      "Batch size: 32\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.0001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: None\n",
      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x7ff5195e7520>\n",
      "Monitor: pr_auc\n",
      "Monitor criterion: max\n",
      "Epochs: 15\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 / 15: 100%|██████████| 243/243 [00:03<00:00, 66.17it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-0, step-243 ---\n",
      "loss: 1.2175\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 227.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-0, step-243 ---\n",
      "pr_auc: 0.6090\n",
      "roc_auc: 0.5884\n",
      "f1: 0.6103\n",
      "loss: 0.7941\n",
      "New best pr_auc score (0.6090) at epoch-0, step-243\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1 / 15: 100%|██████████| 243/243 [00:02<00:00, 95.22it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-1, step-486 ---\n",
      "loss: 1.0263\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 270.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-1, step-486 ---\n",
      "pr_auc: 0.6337\n",
      "roc_auc: 0.6155\n",
      "f1: 0.6277\n",
      "loss: 0.7491\n",
      "New best pr_auc score (0.6337) at epoch-1, step-486\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2 / 15: 100%|██████████| 243/243 [00:02<00:00, 93.65it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-2, step-729 ---\n",
      "loss: 0.9085\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 267.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-2, step-729 ---\n",
      "pr_auc: 0.6457\n",
      "roc_auc: 0.6255\n",
      "f1: 0.6238\n",
      "loss: 0.7093\n",
      "New best pr_auc score (0.6457) at epoch-2, step-729\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3 / 15: 100%|██████████| 243/243 [00:03<00:00, 76.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-3, step-972 ---\n",
      "loss: 0.8412\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 198.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-3, step-972 ---\n",
      "pr_auc: 0.6532\n",
      "roc_auc: 0.6369\n",
      "f1: 0.6521\n",
      "loss: 0.7020\n",
      "New best pr_auc score (0.6532) at epoch-3, step-972\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 4 / 15: 100%|██████████| 243/243 [00:02<00:00, 91.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-4, step-1215 ---\n",
      "loss: 0.7811\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 276.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-4, step-1215 ---\n",
      "pr_auc: 0.6603\n",
      "roc_auc: 0.6385\n",
      "f1: 0.6323\n",
      "loss: 0.6859\n",
      "New best pr_auc score (0.6603) at epoch-4, step-1215\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 5 / 15: 100%|██████████| 243/243 [00:02<00:00, 99.31it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-5, step-1458 ---\n",
      "loss: 0.7451\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 270.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-5, step-1458 ---\n",
      "pr_auc: 0.6580\n",
      "roc_auc: 0.6382\n",
      "f1: 0.6291\n",
      "loss: 0.6849\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 6 / 15: 100%|██████████| 243/243 [00:02<00:00, 93.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-6, step-1701 ---\n",
      "loss: 0.7158\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 224.82it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-6, step-1701 ---\n",
      "pr_auc: 0.6562\n",
      "roc_auc: 0.6378\n",
      "f1: 0.6371\n",
      "loss: 0.6838\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 7 / 15: 100%|██████████| 243/243 [00:02<00:00, 90.81it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-7, step-1944 ---\n",
      "loss: 0.6777\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 278.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-7, step-1944 ---\n",
      "pr_auc: 0.6590\n",
      "roc_auc: 0.6417\n",
      "f1: 0.6172\n",
      "loss: 0.6782\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 8 / 15: 100%|██████████| 243/243 [00:02<00:00, 95.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-8, step-2187 ---\n",
      "loss: 0.6581\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 281.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-8, step-2187 ---\n",
      "pr_auc: 0.6592\n",
      "roc_auc: 0.6462\n",
      "f1: 0.6437\n",
      "loss: 0.6810\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 9 / 15: 100%|██████████| 243/243 [00:03<00:00, 73.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-9, step-2430 ---\n",
      "loss: 0.6387\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 272.67it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-9, step-2430 ---\n",
      "pr_auc: 0.6576\n",
      "roc_auc: 0.6428\n",
      "f1: 0.6589\n",
      "loss: 0.6868\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 10 / 15: 100%|██████████| 243/243 [00:02<00:00, 84.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-10, step-2673 ---\n",
      "loss: 0.6261\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 267.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-10, step-2673 ---\n",
      "pr_auc: 0.6570\n",
      "roc_auc: 0.6419\n",
      "f1: 0.6217\n",
      "loss: 0.6812\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 11 / 15: 100%|██████████| 243/243 [00:03<00:00, 73.52it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-11, step-2916 ---\n",
      "loss: 0.6116\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 276.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-11, step-2916 ---\n",
      "pr_auc: 0.6522\n",
      "roc_auc: 0.6367\n",
      "f1: 0.5841\n",
      "loss: 0.6851\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 12 / 15: 100%|██████████| 243/243 [00:02<00:00, 92.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-12, step-3159 ---\n",
      "loss: 0.6066\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 278.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-12, step-3159 ---\n",
      "pr_auc: 0.6551\n",
      "roc_auc: 0.6330\n",
      "f1: 0.5935\n",
      "loss: 0.6867\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 13 / 15: 100%|██████████| 243/243 [00:03<00:00, 77.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-13, step-3402 ---\n",
      "loss: 0.5941\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 226.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-13, step-3402 ---\n",
      "pr_auc: 0.6570\n",
      "roc_auc: 0.6378\n",
      "f1: 0.6110\n",
      "loss: 0.6855\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 14 / 15: 100%|██████████| 243/243 [00:02<00:00, 90.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-14, step-3645 ---\n",
      "loss: 0.5834\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 31/31 [00:00<00:00, 269.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-14, step-3645 ---\n",
      "pr_auc: 0.6611\n",
      "roc_auc: 0.6421\n",
      "f1: 0.6327\n",
      "loss: 0.6925\n",
      "New best pr_auc score (0.6611) at epoch-14, step-3645\n",
      "Loaded best model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 32/32 [00:00<00:00, 264.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'pr_auc': 0.697333603874223, 'roc_auc': 0.6722280364429949, 'f1': 0.6672694394213382, 'loss': 0.6479178862646222}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# STEP 4: define trainer\n",
    "trainer = Trainer(model=model_no_pre)\n",
    "trainer.train(\n",
    "    train_dataloader=train_dataloader,\n",
    "    val_dataloader=val_dataloader,\n",
    "    epochs=15,\n",
    "    optimizer_params = {\"lr\": 1e-4},\n",
    "    monitor=\"pr_auc\",\n",
    ")\n",
    "\n",
    "# STEP 5: evaluate\n",
    "print(trainer.evaluate(test_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.16 ('txgnn_env')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "79cb95e61c4f960f4e102f21c45668d32cb5c494b237694c15d64b50342e6e99"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
